{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "from arguments import get_config\n",
    "from interfaces import setup_LMP\n",
    "from visualizers import ValueMapVisualizer\n",
    "from envs.rlbench_env import VoxPoserRLBench\n",
    "from utils import set_lmp_objects\n",
    "import numpy as np\n",
    "from rlbench import tasks\n",
    "\n",
    "openai.api_key = None  # set your API key here"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = get_config('rlbench')\n",
    "# uncomment this if you'd like to change the language model (e.g., for faster speed or lower cost)\n",
    "# for lmp_name, cfg in config['lmp_config']['lmps'].items():\n",
    "#     cfg['model'] = 'gpt-3.5-turbo'\n",
    "\n",
    "# initialize env and voxposer ui\n",
    "visualizer = ValueMapVisualizer(config['visualizer'])\n",
    "env = VoxPoserRLBench(visualizer=visualizer)\n",
    "lmps, lmp_env = setup_LMP(env, config, debug=False)\n",
    "voxposer_ui = lmps['plan_ui']"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Playground\n",
    "\n",
    "By default we use one of the instructions that come with each task. However, you may treat each task as simply a scene initialization from RLBench, and feel free to try any task you can come up with for the scene.\n",
    "\n",
    "Note:\n",
    "- Whether an instruction can be executed or not depends on 1) whether relevant objects are available, and 2) capabilities of the overall algorithm.\n",
    "- Each execution may produce one or more visualizations. You may view them in \"./visualizations/\" folder.\n",
    "- The prompts are adapted with minimal change from the real-world environment in the VoxPoser paper. If a task failure is due to incorrect code generated by the LLM, feel free to modify the relevant prompt in \"./prompts/\" folder.\n",
    "- You may view the reward by printing \"env.latest_reward\". These are computed by RLBench for each task.\n",
    "- To inspect in viewer without performing any action, you may call \"env.rlbench_env._scene.step()\" in a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # uncomment this to show all available tasks in rlbench\n",
    "# # NOTE: in order to run a new task, you need to add the list of objects (and their corresponding env names) to the \"task_object_names.json\" file. See README for more details.\n",
    "# print([task for task in dir(tasks) if task[0].isupper() and not '_' in task])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# below are the tasks that have object names added to the \"task_object_names.json\" file\n",
    "# uncomment one to use\n",
    "env.load_task(tasks.PutRubbishInBin)\n",
    "# env.load_task(tasks.LampOff)\n",
    "# env.load_task(tasks.OpenWineBottle)\n",
    "# env.load_task(tasks.PushButton)\n",
    "# env.load_task(tasks.TakeOffWeighingScales)\n",
    "# env.load_task(tasks.MeatOffGrill)\n",
    "# env.load_task(tasks.SlideBlockToTarget)\n",
    "# env.load_task(tasks.TakeLidOffSaucepan)\n",
    "# env.load_task(tasks.TakeUmbrellaOutOfUmbrellaStand)\n",
    "\n",
    "descriptions, obs = env.reset()\n",
    "set_lmp_objects(lmps, env.get_object_names())  # set the object names to be used by voxposer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instruction = np.random.choice(descriptions)\n",
    "voxposer_ui(instruction)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "voxposer-rlbench",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "d3e8e1ce6146f4dbb0e1ef1d424d742293054e718434205b504bd28714852756"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
